Rejection Sampling Variational Inference

Variational inference using the reparameterization trick has enabled large-scale approximate Bayesian inference in complex probabilistic models, leveraging stochastic optimization to sidestep intractable expectations. The reparameterization trick is applicable when we can simulate a random variable by applying a (differentiable) deterministic function on an auxiliary random variable whose distribution is fixed. For many distributions of interest (such as the gamma or Dirichlet), simulation of random variables relies on rejection sampling. The discontinuity introduced by the accept–reject step means that standard reparameterization tricks are not applicable. We propose a new method that lets us leverage reparameterization gradients even when variables are outputs of a rejection sampling algorithm. Our approach enables reparameterization on a larger class of variational distributions. In several studies of real and synthetic data, we show that the variance of the estimator of the gradient is significantly lower than other state-of-the-art methods. This leads to faster convergence of stochastic optimization variational inference. Let p(x, z) be a probabilistic model, i.e., a joint probability distribution of data x and latent (unobserved) variables z. In Bayesian inference, we are interested in the posterior distribution p(z|x) = p(x,z) p(x) . For most models, the posterior distribution is analytically intractable and we have to use an approximation, such as Monte Carlo methods or variational inference. In this paper, we focus on variational inference. In variational inference, we approximate the posterior with a variational family of distributions q(z ; θ), parameterized by θ. Typically, we choose the variational parameters θ that minimize the Kullback-Leibler (KL) divergence between q(z ; θ) and p(z|x). This minimization is equivalent to maximizing the evidence lower bound (ELBO) [Jordan et al., 1999], L(θ) = Eq(z ;θ) [f(z)] +H[q(z ; θ)], f(z) := log p(x, z), H[q(z ; θ)] := Eq(z ;θ)[− log q(z ; θ)]. (1) When the model and variational family satisfy conjugacy requirements, we can use coordinate ascent to find a local optimum of the ELBO [Ghahramani and Beal, 2001, Blei et al., 2016]. If the conjugacy requirements are not satisfied, a common approach is to build a Monte Carlo estimator of the gradient of the ELBO [Paisley et al., 2012, Ranganath et al., 2014, Salimans and Knowles, 2013]. Empirically, the reparameterization trick has been shown to be beneficial over direct Monte Carlo estimation of the gradient using the score fuction estimator [Rezende et al., 2014, Kingma and Welling, 2014, Titsias and Lázaro-Gredilla, 2014, Fan et al., 2015]. However, it is not generally applicable, it requires that: (i) the latent variables z are continuous; and (ii) we can simulate from q(z ; θ) as follows, z = h(ε, θ), with ε ∼ s(ε). (2) Here, s(ε) is a distribution that does not depend on the variational parameters; it is typically a standard normal or a standard uniform. Further, h(ε, θ) is differentiable with respect to θ. Using (2), we can move the derivative inside the expectation and rewrite the gradient of the ELBO as ∇θL(θ) = Es(ε) [∇zf(h(ε, θ))∇θh(ε, θ)] +∇θH[q(z ; θ)]. ∗Corresponding author: [email protected] Advances in Approximate Bayesian Inference (NIPS 2016 Workshop), Barcelona, Spain. Algorithm 1 Reparameterized Rejection Sampling Input: target q(z ; θ), proposal r(z ; θ), and constant Mθ, with q(z ; θ) ≤Mθr(z ; θ) Output: ε such that h(ε, θ) ∼ q(z ; θ) 1: i← 0 2: repeat 3: i← i+ 1 4: Propose εi ∼ s(ε) 5: Simulate ui ∼ U [0, 1] 6: until ui < q(h(εi,θ) ;θ) Mθr(h(εi,θ) ;θ) 7: return εi 6 4 2 0 2 4 6 ε 10 10 10 10 10 10 10 10 10